import cv2
import numpy as np
import socket
import pickle
import struct
import time
from collections import defaultdict

class ScreenClient:
    def __init__(self, listen_port=5005):
        # UDP配置
        self.udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.udp_socket.bind(("", listen_port))
        self.udp_socket.settimeout(1.0)
        
        # 屏幕参数
        self.width = 0
        self.height = 0
        self.grid_rows = 4
        self.grid_cols = 4
        self.block_width = 0
        self.block_height = 0
        
        # 内存存储（不写入磁盘）
        self.received_blocks = {}  # (i,j) -> 解压后的块图像
        self.packet_buffer = defaultdict(dict)  # 分片重组缓冲区
        self.running = True

    def receive_and_reassemble(self):
        """接收并重组分片数据"""
        try:
            while True:
                packet, _ = self.udp_socket.recvfrom(4096)
                if len(packet) < 16:
                    continue
                
                # 解析包头
                header = packet[:16]
                data = packet[16:]
                packet_id, total_packets, current_packet, chunk_size = struct.unpack(">IIII", header)
                
                if len(data) != chunk_size:
                    continue
                
                # 存储分片
                self.packet_buffer[packet_id][current_packet] = data
                
                # 检查是否收全
                if len(self.packet_buffer[packet_id]) == total_packets:
                    full_data = b"".join([self.packet_buffer[packet_id][i] for i in range(total_packets)])
                    del self.packet_buffer[packet_id]
                    return pickle.loads(full_data)
        except (socket.timeout, EOFError):
            return None

    def decompress_block(self, compressed_data):
        """解压JPEG数据（仅内存操作）"""
        if not compressed_data:
            return None
        nparr = np.frombuffer(compressed_data, np.uint8)
        return cv2.imdecode(nparr, cv2.IMREAD_COLOR)

    def init_screen_params(self, width, height):
        """初始化屏幕参数"""
        self.width = width
        self.height = height
        self.block_width = self.width // self.grid_cols
        self.block_height = self.height // self.grid_rows
        print(f"初始化屏幕: {width}x{height}")

    def update_blocks(self, changed_blocks):
        """更新内存中的块数据"""
        for i, j, compressed_block in changed_blocks:
            block = self.decompress_block(compressed_block)
            if block is not None:
                self.received_blocks[(i, j)] = block.copy()

    def draw_mouse(self, frame, mouse_pos):
        """绘制鼠标指针（内存中绘制）"""
        x, y = mouse_pos
        cursor_size = 20
        cursor = np.zeros((cursor_size, cursor_size, 3), dtype=np.uint8)
        cv2.circle(cursor, (10, 10), 5, (255, 255, 255), -1)  # 白色圆点
        cv2.circle(cursor, (10, 10), 6, (0, 0, 255), 1)       # 红色边框
        
        draw_x = x - 10
        draw_y = y - 10
        
        start_y = max(0, draw_y)
        end_y = min(frame.shape[0], draw_y + cursor_size)
        start_x = max(0, draw_x)
        end_x = min(frame.shape[1], draw_x + cursor_size)
        
        # 绘制鼠标到帧上
        frame[start_y:end_y, start_x:end_x] = cursor[
            max(0, -draw_y):max(0, -draw_y)+(end_y-start_y),
            max(0, -draw_x):max(0, -draw_x)+(end_x-start_x)
        ]
        return frame

    def stitch_frame(self, mouse_pos):
        """拼接内存中的块并绘制鼠标"""
        if not self.received_blocks or self.width == 0:
            return None
        
        # 内存中拼接块
        stitched = np.zeros((self.height, self.width, 3), dtype=np.uint8)
        for (i, j), block in self.received_blocks.items():
            y1 = i * self.block_height
            y2 = (i + 1) * self.block_height if i < self.grid_rows - 1 else self.height
            x1 = j * self.block_width
            x2 = (j + 1) * self.block_width if j < self.grid_cols - 1 else self.width
            
            block = cv2.resize(block, (x2 - x1, y2 - y1))
            stitched[y1:y2, x1:x2] = block
        
        # 绘制鼠标
        return self.draw_mouse(stitched, mouse_pos)

    def run(self):
        print("客户端启动，等待数据...")
        try:
            while self.running:
                data = self.receive_and_reassemble()
                if not data:
                    continue
                
                if "screen_info" in data:
                    self.init_screen_params(*data["screen_info"])
                
                if "changed_blocks" in data and data["changed_blocks"]:
                    self.update_blocks(data["changed_blocks"])
                
                # 显示内存中的拼接结果
                frame = self.stitch_frame(data["mouse"])
                if frame is not None:
                    cv2.imshow("Screen Stream (Press 'q' to exit)", frame)
                
                if cv2.waitKey(1) & 0xFF == ord('q'):
                    self.running = False
                
                time.sleep(0.01)
        except KeyboardInterrupt:
            print("客户端停止")
        finally:
            self.udp_socket.close()
            cv2.destroyAllWindows()

if __name__ == "__main__":
    client = ScreenClient()
    client.run()